from .base_exchange import *
from bitget import Bitget
from bitget.ws import BitgetWsClient
import asyncio
import time

class BitgetExchange(BaseExchange):
    def __init__(self, api_key: str = "", api_secret: str = "", passphrase: str = "", 
                 is_futures: bool = False):
        super().__init__(api_key, api_secret, passphrase)
        self.name = "Bitget"
        self.is_futures = is_futures
        
        # 初始化REST API客户端
        self.client = Bitget(
            api_key=api_key,
            api_secret=api_secret,
            passphrase=passphrase
        )
        
        # WebSocket客户端
        self.ws = None
        self.ws_connected = False
        
        # 缓存数据
        self.tickers = {}
        self.orderbooks = {}
        self.balances = {}
        self.positions = {}
        
    async def connect_ws(self):
        """连接WebSocket"""
        try:
            def handle_message(message):
                asyncio.create_task(self._process_ws_message(message))
                
            self.ws = BitgetWsClient(
                api_key=self.api_key,
                api_secret=self.api_secret,
                passphrase=self.passphrase,
                use_server_time=True,
                recv_window=5000
            )
            
            # 设置回调
            self.ws.on_message = handle_message
            
            # 连接WebSocket
            await self.ws.connect()
            
            # 订阅频道
            await self._subscribe_channels()
            
            self.ws_connected = True
            
        except Exception as e:
            raise Exception(f"Failed to connect WebSocket: {str(e)}")
            
    async def _subscribe_channels(self):
        """订阅WebSocket频道"""
        try:
            # 订阅行情数据
            instType = "mc" if self.is_futures else "sp"
            channels = [
                {
                    "instType": instType,
                    "channel": "ticker",
                    "instId": "BTCUSDT"
                },
                {
                    "instType": instType,
                    "channel": "depth",
                    "instId": "BTCUSDT"
                }
            ]
            
            if self.api_key:
                # 订阅账户和持仓数据
                channels.extend([
                    {
                        "instType": instType,
                        "channel": "account"
                    },
                    {
                        "instType": instType,
                        "channel": "positions"
                    }
                ])
                
            await self.ws.subscribe(channels)
            
        except Exception as e:
            raise Exception(f"Failed to subscribe channels: {str(e)}")
            
    async def _process_ws_message(self, message):
        """处理WebSocket消息"""
        try:
            if not isinstance(message, dict):
                return
                
            channel = message.get('arg', {}).get('channel')
            if not channel:
                return
                
            if channel == 'ticker':
                await self._handle_ticker(message)
            elif channel == 'depth':
                await self._handle_orderbook(message)
            elif channel == 'account':
                await self._handle_account(message)
            elif channel == 'positions':
                await self._handle_position(message)
                
        except Exception as e:
            print(f"Error processing message: {e}")
            
    async def _handle_ticker(self, message):
        """处理行情数据"""
        try:
            data = message['data']
            symbol = data['instId']
            self.tickers[symbol] = {
                'last': float(data['last']),
                'bid': float(data['bestBid']),
                'ask': float(data['bestAsk']),
                'volume': float(data['baseVolume']),
                'timestamp': int(data['ts'])
            }
        except Exception as e:
            print(f"Error handling ticker: {e}")
            
    async def get_ticker(self, symbol: str) -> Ticker:
        """获取最新行情"""
        try:
            if self.is_futures:
                response = self.client.mix_market_api().get_ticker(symbol)
            else:
                response = self.client.spot_market_api().get_ticker(symbol)
                
            data = response['data']
            return Ticker(
                symbol=symbol,
                last_price=float(data['last']),
                bid_price=float(data['bestBid']),
                ask_price=float(data['bestAsk']),
                volume_24h=float(data['baseVolume']),
                timestamp=int(data['ts'])
            )
        except Exception as e:
            raise Exception(f"Failed to get ticker: {str(e)}")
            
    async def get_orderbook(self, symbol: str, limit: int = 20) -> OrderBook:
        """获取订单簿"""
        try:
            if self.is_futures:
                response = self.client.mix_market_api().get_depth(symbol, limit=limit)
            else:
                response = self.client.spot_market_api().get_depth(symbol, limit=limit)
                
            data = response['data']
            return OrderBook(
                asks=[(float(price), float(qty)) for price, qty in data['asks']],
                bids=[(float(price), float(qty)) for price, qty in data['bids']],
                timestamp=int(data['ts'])
            )
        except Exception as e:
            raise Exception(f"Failed to get orderbook: {str(e)}")
            
    async def get_balance(self, asset: str = "") -> Dict[str, Balance]:
        """获取账户余额"""
        try:
            balances = {}
            if self.is_futures:
                response = self.client.mix_account_api().get_account(symbol="USDT")
            else:
                response = self.client.spot_account_api().get_account(asset)
                
            for item in response['data']:
                ccy = item['coin']
                if not asset or ccy == asset:
                    balances[ccy] = Balance(
                        asset=ccy,
                        free=float(item['available']),
                        locked=float(item['frozen']),
                        total=float(item['total']),
                        timestamp=time.time()
                    )
            return balances
        except Exception as e:
            raise Exception(f"Failed to get balance: {str(e)}")
            
    async def get_positions(self, symbol: str = "") -> Dict[str, Position]:
        """获取持仓信息"""
        try:
            if not self.is_futures:
                return {}
                
            response = self.client.mix_account_api().get_positions(
                symbol=symbol if symbol else None
            )
            
            positions = {}
            for pos in response['data']:
                symbol = pos['symbol']
                positions[symbol] = Position(
                    symbol=symbol,
                    size=float(pos['total']),
                    entry_price=float(pos['averageOpenPrice']),
                    leverage=float(pos['leverage']),
                    margin=float(pos['margin']),
                    unrealized_pnl=float(pos['unrealizedPnl']),
                    liquidation_price=float(pos['liquidationPrice']),
                    timestamp=time.time()
                )
            return positions
        except Exception as e:
            raise Exception(f"Failed to get positions: {str(e)}")
            
    async def create_market_order(self, symbol: str, side: OrderSide, amount: float) -> Order:
        """创建市价单"""
        try:
            if self.is_futures:
                response = self.client.mix_order_api().place_order(
                    symbol=symbol,
                    marginCoin='USDT',
                    size=str(amount),
                    side=side.value,
                    orderType='market',
                    timeInForceValue='normal'
                )
            else:
                response = self.client.spot_order_api().place_order(
                    symbol=symbol,
                    side=side.value,
                    orderType='market',
                    size=str(amount)
                )
                
            if response['code'] != '00000':
                raise Exception(response['msg'])
                
            order_id = response['data']['orderId']
            return await self.get_order(symbol, order_id)
            
        except Exception as e:
            raise Exception(f"Failed to create market order: {str(e)}")
            
    async def create_limit_order(self, symbol: str, side: OrderSide, 
                               price: float, amount: float) -> Order:
        """创建限价单"""
        try:
            if self.is_futures:
                response = self.client.mix_order_api().place_order(
                    symbol=symbol,
                    marginCoin='USDT',
                    size=str(amount),
                    price=str(price),
                    side=side.value,
                    orderType='limit',
                    timeInForceValue='post_only'
                )
            else:
                response = self.client.spot_order_api().place_order(
                    symbol=symbol,
                    side=side.value,
                    orderType='limit',
                    price=str(price),
                    size=str(amount)
                )
                
            if response['code'] != '00000':
                raise Exception(response['msg'])
                
            order_id = response['data']['orderId']
            return await self.get_order(symbol, order_id)
            
        except Exception as e:
            raise Exception(f"Failed to create limit order: {str(e)}")
            
    async def set_leverage(self, symbol: str, leverage: int) -> bool:
        """设置杠杆倍数"""
        try:
            if not self.is_futures:
                return False
                
            response = self.client.mix_account_api().set_leverage(
                symbol=symbol,
                marginCoin='USDT',
                leverage=leverage
            )
            
            return response['code'] == '00000'
            
        except Exception as e:
            raise Exception(f"Failed to set leverage: {str(e)}")
            
    async def get_funding_rate(self, symbol: str) -> float:
        """获取资金费率"""
        try:
            if not self.is_futures:
                return 0.0
                
            response = self.client.mix_market_api().get_funding_rate(symbol)
            
            if response['code'] != '00000':
                raise Exception(response['msg'])
                
            return float(response['data']['fundingRate'])
            
        except Exception as e:
            raise Exception(f"Failed to get funding rate: {str(e)}")
            
    async def get_funding_time(self, symbol: str) -> int:
        """获取下次资金费时间"""
        try:
            if not self.is_futures:
                return 0
                
            response = self.client.mix_market_api().get_funding_time(symbol)
            
            if response['code'] != '00000':
                raise Exception(response['msg'])
                
            return int(response['data']['nextFundingTime'])
            
        except Exception as e:
            raise Exception(f"Failed to get funding time: {str(e)}")
            
    async def get_order(self, symbol: str, order_id: str) -> Order:
        """获取订单信息"""
        try:
            if self.is_futures:
                response = self.client.mix_order_api().get_order_detail(
                    symbol=symbol,
                    orderId=order_id
                )
            else:
                response = self.client.spot_order_api().get_order_detail(
                    symbol=symbol,
                    orderId=order_id
                )
                
            if response['code'] != '00000':
                raise Exception(response['msg'])
                
            data = response['data']
            return Order(
                exchange_id=self.name,
                order_id=order_id,
                symbol=symbol,
                order_type=OrderType.MARKET if data['orderType'] == 'market' else OrderType.LIMIT,
                side=OrderSide.BUY if data['side'] == 'buy' else OrderSide.SELL,
                price=float(data['price']),
                amount=float(data['size']),
                filled=float(data['filledQty']),
                status=self._convert_status(data['status']),
                timestamp=int(data['cTime'])
            )
            
        except Exception as e:
            raise Exception(f"Failed to get order: {str(e)}")
            
    def _convert_status(self, status: str) -> OrderStatus:
        """转换订单状态"""
        status_map = {
            'new': OrderStatus.PENDING,
            'live': OrderStatus.OPEN,
            'filled': OrderStatus.FILLED,
            'cancelled': OrderStatus.CANCELED,
            'rejected': OrderStatus.REJECTED
        }
        return status_map.get(status.lower(), OrderStatus.PENDING) 